package rate

import (
	"context"
	"sync/atomic"

	"go-toolkit/library/net/metadata"

	"fmt"

	"golang.org/x/time/rate"
	"google.golang.org/grpc"
)

type limiterMap map[string]*rate.Limiter

// Rule 定义限流规则
type Rule interface {
	// 返回规则名
	Name() string
	// 从请求中抽取 limiterMap 的 key，抽不出返回空字符串
	LimitKey(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo) string
}

type ruleBucket struct {
	r      Rule
	limMap atomic.Value // limiterMap
}

func newRuleBucket(r Rule) *ruleBucket {
	return &ruleBucket{r: r}
}

func (rb *ruleBucket) Reload(conf Config) {
	olm, _ := rb.limMap.Load().(limiterMap)

	lm := make(limiterMap, len(conf))
	for _, v := range conf {
		k := v.Key

		if l, ok := olm[k]; ok && l.Limit() == v.Limit && l.Burst() == v.Burst {
			lm[k] = l
		} else {
			lm[k] = rate.NewLimiter(v.fix())
		}
	}

	rb.limMap.Store(lm)
}

// Allow 返回 deny 的 key 名
func (rb *ruleBucket) Allow(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo) (allow bool, key string) {
	allow = true
	key = rb.r.LimitKey(ctx, req, args)
	if key == "" {
		return
	}

	lm, _ := rb.limMap.Load().(limiterMap)
	if l, ok := lm[key]; ok {
		allow = l.Allow()
	}

	return
}

// RuleCaller 以 caller 限流的规则
type RuleCaller struct{}

// Name 规则名
func (r RuleCaller) Name() string {
	return "caller"
}

// LimitKey 组合 key
func (r RuleCaller) LimitKey(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo) string {
	caller := metadata.String(ctx, metadata.Caller)
	if caller == "" {
		caller = "no_user"
	}
	return caller
}

// RuleMethod 以 method 限流的规则
type RuleMethod struct{}

// Name 规则名
func (r RuleMethod) Name() string {
	return "method"
}

// LimitKey 组合 key
func (r RuleMethod) LimitKey(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo) string {
	return args.FullMethod
}

// RuleMethodCaller 组合 method caller 限流的规则
type RuleMethodCaller struct {
	rc RuleCaller
	rm RuleMethod
}

// Name 规则名
func (r RuleMethodCaller) Name() string {
	return "method_caller"
}

// LimitKey 组合 key
func (r RuleMethodCaller) LimitKey(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo) string {
	rmKey := r.rm.LimitKey(ctx, req, args)
	if rmKey == "" {
		return ""
	}

	return fmt.Sprintf("%s|%s", rmKey, r.rc.LimitKey(ctx, req, args))
}
